import torch.nn as nn
import numpy as np
import torch


class ResidualBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, downsample=None):
        super().__init__()
        self.conv1 = nn.Conv1d(in_channels, out_channels,kernel_size, stride, padding=kernel_size//2, bias=False)
        self.bn1 = nn.BatchNorm1d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv1d(out_channels, out_channels,kernel_size, stride, padding=kernel_size//2, bias=False)
        self.bn2 = nn.BatchNorm1d(out_channels)
        self.downsample = downsample

    def forward(self, x):
        indentity = x
        if self.downsample is not None:
            indentity = self.downsample(x)
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        out += indentity
        out = self.relu(out)
        return out


class ResNet1D(nn.Module):
    def __init__(self, num_class=2, embed_audio=None):  # 2分类
        super().__init__()
        self.layer1 = nn.Sequential(
            nn.Conv1d(1, 64, kernel_size=7, stride=2, padding=3, bias=False),  # [32, 64, 150]
            nn.BatchNorm1d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool1d(kernel_size=3, stride=2, padding=1)  # [32, 64, 75]
        )
        self.layer2 = self._make_layer(64, 128, 2)            # []
        self.layer3 = self._make_layer(128, 256, 2)           # []
        self.layer4 = self._make_layer(256, 256, 2)
        self.layer5 = self._make_layer(512, 512, 2)
        # self.layerl_ese = self._make_layer(256, 512, 2)
        self.avg_pool = nn.AdaptiveAvgPool1d(1)
        self.fc = nn.Linear(512, num_class)

        # 定义音频特征卷积
        self.layer6 = nn.Sequential(
            nn.Conv1d(1, 256, kernel_size=3, stride=2, padding=1),  # 输出形状: [32, 256, 256]
            nn.Conv1d(256, 256, kernel_size=3, stride=3, padding=1),  # 输出形状: [32, 256, 86]
            nn.AdaptiveAvgPool1d(75)                              # 自适应池化到目标形状: [32, 256, 75]
        )
    
    def _make_layer(self, in_channels, out_channels, blocks):
        downsample = nn.Sequential(
            nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False),
            nn.BatchNorm1d(out_channels)
        )
        layers = [ResidualBlock(in_channels, out_channels, kernel_size=3 ,downsample=downsample)]
        for _ in range(1, blocks):
            layers.append(ResidualBlock(out_channels, out_channels, kernel_size=3))
        return nn.Sequential(*layers)
    
    def forward(self, txt_x, embed_y):
        x = self.layer1(txt_x)  # [32, 64, 75]
        x = self.layer2(x)      # []
        x = self.layer3(x)
        x = self.layer4(x)
        # print(x.shape)
        embed_y = self.layer6(embed_y)   
        # print(embed_y.shape)
        x = torch.cat((x, embed_y), dim=1)
        # print(x.shape)
        # x = self.layerl_ese(embed_y)       # 这是只有音频的输入情况
        x = self.layer5(x)
        x = self.avg_pool(x).squeeze(-1)
        x = self.fc(x)
        return x
